import numpy as np
from math import *
import cv2 as cv
import time

class room:
    free_id = 1

    def __init__(self):
        self.connected_rooms: 'list[room]' = [None] * 4
        self.backwards_doors: 'list[int]' = [None] * 4
        self.id = room.free_id
        room.free_id += 1
    
    def disconnect(self, door: int):
        self.connected_rooms[door].connected_rooms[self.backwards_doors[door]] = None
        self.connected_rooms[door].backwards_doors[self.backwards_doors[door]] = None
        self.connected_rooms[door] = None
        self.backwards_doors[door] = None

    @staticmethod
    def connect(room1: 'room', door1to2: int, room2: 'room', door2to1: int):
        if not room1.connected_rooms[door1to2] is None:
            room1.disconnect(door1to2)
        if not room2.connected_rooms[door2to1] is None:
            room2.disconnect(door2to1)
        room1.connected_rooms[door1to2] = room2
        room1.backwards_doors[door1to2] = door2to1
        room2.connected_rooms[door2to1] = room1
        room2.backwards_doors[door2to1] = door1to2
    
    def __eq__(self, other: 'room') -> bool:
        return self.id == other.id

def rotate(vec2: np.ndarray, angle: float) -> np.ndarray:
    return np.array([vec2[0] * cos(angle) - vec2[1] * sin(angle), vec2[1] * cos(angle) + vec2[0] * sin(angle)])

def cross(vec1: np.ndarray, vec2: np.ndarray) -> float:
    return vec1[0] * vec2[1] - vec1[1] * vec2[0]

def get_intersection(org1: np.ndarray, vec1: np.ndarray, org2: np.ndarray, vec2: np.ndarray) -> float:
    if cross(vec1, vec2) == 0: return float('inf')
    return cross(org2 - org1, vec2) / cross(vec1, vec2)

def slice_polygon(polygon: 'list[np.ndarray]', origin: np.ndarray, dir: np.ndarray) -> 'tuple[list[np.ndarray], list[np.ndarray]]':
    faces = []
    points = []
    for i in range(len(polygon)):
        t = get_intersection(polygon[i - 1], polygon[i] - polygon[i - 1], origin, dir)
        point = polygon[i - 1] * (1 - t) + polygon[i] * t
        if 0 <= t <= 1 and (not points or not (point[None, :] == np.array(points)).all(axis=1).any()):
            faces.append(i)
            points.append(point)
            if len(faces) == 2:
                break
    if len(faces) != 2: return [polygon]
    p1 = polygon[:faces[0]] + points + polygon[faces[1]:]
    p2 = polygon[faces[0]: faces[1]] + points[::-1]
    return [p1, p2]

class backrooms:
    def __init__(self, rooms: 'list[room]' = [], origin_room: int = 0,
                 origin_pos: 'tuple[float]' = (.5, .5), rotation: float = 0.,
                 room_size: int = 100, image_size: 'tuple[int]' = (1080, 720)):
        self.rooms = rooms
        self.origin_room = origin_room
        self.origin_pos = np.array(origin_pos)
        self.rotation = rotation
        self.room_size = room_size
        self.image_size = np.array(image_size)

    def divide_image(self):
        min_room_coords = np.array([100000, 100000])
        max_room_coords = np.array([-100000, -100000])
        for corner in np.array([[1, -1], [1, 1], [-1, 1], [-1, -1]]):
            corner_coords = rotate(self.image_size / 2 / self.room_size * corner, -self.rotation) + self.origin_pos
            min_room_coords = np.minimum(min_room_coords, np.floor(corner_coords).astype(int))
            max_room_coords = np.maximum(max_room_coords, -np.floor(-corner_coords).astype(int))
        grid = [[[] for _ in range(max_room_coords[1] - min_room_coords[1])] for _ in range(max_room_coords[0] - min_room_coords[0])]
        for i in range(max_room_coords[0] - min_room_coords[0]):
            for j in range(max_room_coords[1] - min_room_coords[1]):
                org = np.array([i, j]) + min_room_coords
                polygon = []
                for corner in np.array([[1, 0], [1, 1], [0, 1], [0, 0]]):
                    polygon.append(org + corner)
                grid[i][j].append(polygon)
        for i in range(1, max_room_coords[0] - min_room_coords[0]):
            for j in range(1, max_room_coords[1] - min_room_coords[1]):
                p = np.array([i, j], dtype=float) + min_room_coords
                vec = p - self.origin_pos
                while ((min_room_coords < p) & (p < max_room_coords)).all():
                    cell = np.floor(p).astype(int) - (vec < 0) * (p == np.floor(p))
                    i2, j2 = cell - min_room_coords
                    grid[i2][j2] = sum([slice_polygon(poly, p, vec) for poly in grid[i2][j2]], [])
                    walls = cell + (vec > 0)
                    t = (walls - p) / vec
                    p += vec * np.nan_to_num(t, nan=float('inf')).min()
                    if t[0] < t[1]: p[0] = walls[0]
                    else: p[1] = walls[1]
        polygons = []
        for i in range(max_room_coords[0] - min_room_coords[0]):
            for j in range(max_room_coords[1] - min_room_coords[1]):
                for polygon in grid[i][j]:
                    poly = [self.image_size / 2 + rotate((p - self.origin_pos) * self.room_size, self.rotation) for p in polygon]
                    if not ((0 <= np.array(poly)).all(axis=1).any() and (np.array(poly) <= self.image_size[None, :]).all(axis=1).any()):
                        continue
                    polygons.append({
                        'room': np.array([i, j]) + min_room_coords,
                        'polygon': np.array(poly).astype(int)
                    })
        return polygons
    
    def test_draw(self):
        image = np.zeros(list(self.image_size[::-1]) + [3], dtype='uint8')
        t = time.time()
        polygons = self.divide_image()
        print(time.time() - t)
        for d in polygons:
            color = tuple(map(int, np.random.randint(256, size=3)))
            # if (d['room'] == [-4, 4]).all():
            #     color = (0, 0, 0)
            cv.fillPoly(image, [d['polygon']], color)
        cv.imshow('test', image)
        while 1: cv.waitKey(50)

backrooms(origin_pos=(0.2001, 0.300001), rotation=.5).test_draw()
